import pygame
import random

# 初始化Pygame
pygame.init()

# 颜色定义
COLORS = [
    (0, 0, 0),        # 黑色（背景）
    (120, 195, 239),  # 蓝色
    (165, 105, 189),  # 紫色
    (245, 171, 53),   # 橙色
    (46, 204, 113),   # 绿色
    (231, 76, 60),    # 红色
    (52, 152, 219),   # 深蓝色
    (241, 196, 15)    # 黄色
]

# 俄罗斯方块形状
SHAPES = [
    [[1, 5, 9, 13], [4, 5, 6, 7]],  # I
    [[4, 5, 9, 10], [2, 6, 5, 9]],  # Z
    [[6, 7, 9, 10], [1, 5, 6, 10]], # S
    [[1, 2, 5, 9], [0, 4, 5, 6], [1, 5, 9, 8], [4, 5, 6, 10]], # J
    [[1, 2, 6, 10], [5, 6, 7, 9], [2, 6, 10, 11], [3, 5, 6, 7]], # L
    [[1, 4, 5, 6], [1, 4, 5, 9], [4, 5, 6, 9], [1, 5, 6, 9]], # T
    [[1, 2, 5, 6]], # O
]

# 游戏设置
WINDOW_WIDTH = 500
WINDOW_HEIGHT = 600
BLOCK_SIZE = 30
FIELD_WIDTH = 10
FIELD_HEIGHT = 20
GAME_SPEED = 25

class Tetris:
    def __init__(self, height, width):
        self.height = height
        self.width = width
        self.field = [[0 for _ in range(width)] for _ in range(height)]
        self.score = 0
        self.state = "start"
        self.figure = None
        self.x = 0
        self.y = 0
        self.next_figure = None
        self.level = 1
        
    def new_figure(self):
        if not self.next_figure:
            self.next_figure = Figure(3, 0)
        self.figure = self.next_figure
        self.next_figure = Figure(3, 0)

    def intersects(self):
        intersection = False
        for i in range(4):
            for j in range(4):
                if i * 4 + j in self.figure.image():
                    if (i + self.figure.y > self.height - 1 or
                        j + self.figure.x > self.width - 1 or
                        j + self.figure.x < 0 or
                        self.field[i + self.figure.y][j + self.figure.x] > 0):
                        intersection = True
        return intersection

    def freeze(self):
        for i in range(4):
            for j in range(4):
                if i * 4 + j in self.figure.image():
                    self.field[i + self.figure.y][j + self.figure.x] = self.figure.color
        self.break_lines()
        self.new_figure()
        if self.intersects():
            self.state = "gameover"

    def break_lines(self):
        lines = 0
        for i in range(1, self.height):
            zeros = 0
            for j in range(self.width):
                if self.field[i][j] == 0:
                    zeros += 1
            if zeros == 0:
                lines += 1
                for i1 in range(i, 1, -1):
                    for j in range(self.width):
                        self.field[i1][j] = self.field[i1-1][j]
        self.score += lines ** 2

    def go_space(self):
        while not self.intersects():
            self.figure.y += 1
        self.figure.y -= 1
        self.freeze()

    def go_down(self):
        self.figure.y += 1
        if self.intersects():
            self.figure.y -= 1
            self.freeze()

    def go_side(self, dx):
        old_x = self.figure.x
        self.figure.x += dx
        if self.intersects():
            self.figure.x = old_x

    def rotate(self):
        old_rotation = self.figure.rotation
        self.figure.rotate()
        if self.intersects():
            self.figure.rotation = old_rotation

class Figure:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.type = random.randint(0, len(SHAPES) - 1)
        self.color = random.randint(1, len(COLORS) - 1)
        self.rotation = 0

    def image(self):
        return SHAPES[self.type][self.rotation]

    def rotate(self):
        self.rotation = (self.rotation + 1) % len(SHAPES[self.type])

def main():
    screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
    pygame.display.set_caption("俄罗斯方块")
    clock = pygame.time.Clock()
    game = Tetris(FIELD_HEIGHT, FIELD_WIDTH)
    counter = 0
    pressing_down = False
    done = False

    while not done:
        if game.figure is None:
            game.new_figure()
        counter += 1
        if counter > 100000:
            counter = 0

        if counter % (GAME_SPEED // game.level // 2) == 0 or pressing_down:
            if game.state == "start":
                game.go_down()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                done = True
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    game.rotate()
                if event.key == pygame.K_DOWN:
                    pressing_down = True
                if event.key == pygame.K_LEFT:
                    game.go_side(-1)
                if event.key == pygame.K_RIGHT:
                    game.go_side(1)
                if event.key == pygame.K_SPACE:
                    game.go_space()
                if event.key == pygame.K_ESCAPE:
                    game.__init__(FIELD_HEIGHT, FIELD_WIDTH)

        if event.type == pygame.KEYUP:
            if event.key == pygame.K_DOWN:
                pressing_down = False

        screen.fill((0, 0, 0))

        # 绘制游戏区域
        for i in range(game.height):
            for j in range(game.width):
                pygame.draw.rect(screen, (40, 40, 40),
                               [j * BLOCK_SIZE + 100, i * BLOCK_SIZE,
                                BLOCK_SIZE, BLOCK_SIZE], 1)
                if game.field[i][j] > 0:
                    pygame.draw.rect(screen, COLORS[game.field[i][j]],
                                   [j * BLOCK_SIZE + 100, i * BLOCK_SIZE,
                                    BLOCK_SIZE - 2, BLOCK_SIZE - 1])

        # 绘制当前方块
        if game.figure is not None:
            for i in range(4):
                for j in range(4):
                    p = i * 4 + j
                    if p in game.figure.image():
                        pygame.draw.rect(screen, COLORS[game.figure.color],
                                       [(j + game.figure.x) * BLOCK_SIZE + 100,
                                        (i + game.figure.y) * BLOCK_SIZE,
                                        BLOCK_SIZE - 2, BLOCK_SIZE - 2])

        # 绘制分数
        font = pygame.font.Font(None, 36)
        text = font.render("Score: " + str(game.score), True, (255, 255, 255))
        screen.blit(text, [0, 0])
        text_game_over = font.render("Game Over", True, (255, 255, 255))
        text_press_esc = font.render("Press ESC to restart", True, (255, 255, 255))

        if game.state == "gameover":
            screen.blit(text_game_over, [WINDOW_WIDTH//3, WINDOW_HEIGHT//2])
            screen.blit(text_press_esc, [WINDOW_WIDTH//3 - 30, WINDOW_HEIGHT//2 + 30])

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()

if __name__ == '__main__':
    main()
